import os
from langchain.agents.agent import Agent, AgentExecutor
from langchain_experimental.plan_and_execute import (
    PlanAndExecute,
    load_agent_executor,
    load_chat_planner
)
from langchain_community.chat_models import ChatOpenAI
from langchain.schema import SystemMessage
from langchain.tools import Tool
from langchain.agents import Tool
import requests
import json

# 设置阿里云 DashScope API 密钥
os.environ["DASHSCOPE_API_KEY"] = "你的-dashscope-api-key"


class QwenChatModel:
    """自定义 Qwen 聊天模型类，适配 LangChain 接口"""

    def __init__(self, model_name="qwen-plus", temperature=0.1):
        self.model_name = model_name
        self.temperature = temperature
        self.api_key = 'sk-965dc39b016c49ecbe29de180f4db2b6'
        self.base_url = "https://dashscope.aliyun.com/compatible-mode/v1"

    def invoke(self, messages):
        """调用 Qwen API"""
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        # 转换消息格式
        formatted_messages = []
        for msg in messages:
            if hasattr(msg, 'type'):
                # LangChain 消息对象
                formatted_messages.append({
                    "role": "user" if msg.type == "human" else "assistant",
                    "content": msg.content
                })
            else:
                # 字典格式消息
                role = "user" if msg.get("type") == "human" else "assistant"
                if "role" in msg:
                    role = msg["role"]
                formatted_messages.append({
                    "role": role,
                    "content": msg.get("content", "")
                })

        payload = {
            "model": self.model_name,
            "messages": formatted_messages,
            "temperature": self.temperature,
            "stream": False
        }

        try:
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"]
        except Exception as e:
            print(f"API调用错误: {e}")
            return f"错误: {str(e)}"


# 创建自定义工具
def search_tool(query: str) -> str:
    """搜索工具（示例）"""
    # 这里可以集成真正的搜索API，如 SerperAPI 或 Tavily
    return f"搜索 '{query}' 的结果: 这是模拟的搜索结果，包含相关信息。"


def calculator_tool(expression: str) -> str:
    """计算器工具"""
    try:
        # 简单的表达式计算（生产环境请使用更安全的方法）
        result = eval(expression)
        return f"计算结果: {expression} = {result}"
    except:
        return f"无法计算表达式: {expression}"


def weather_tool(city: str) -> str:
    """天气查询工具（示例）"""
    return f"{city}的天气: 晴天，25°C，湿度60%（模拟数据）"


# 创建工具列表
tools = [
    Tool(
        name="Search",
        func=search_tool,
        description="用于搜索最新信息、新闻或特定主题的知识"
    ),
    Tool(
        name="Calculator",
        func=calculator_tool,
        description="用于执行数学计算，支持基本运算符如 +, -, *, /"
    ),
    Tool(
        name="Weather",
        func=weather_tool,
        description="用于查询城市天气信息"
    )
]


def create_qwen_plan_and_execute_agent():
    """创建基于 Qwen 的 PlanAndExecute 智能体"""

    # 初始化 Qwen 模型
    qwen_model = QwenChatModel(model_name="qwen-plus", temperature=0.1)

    # 创建规划器
    planner = load_chat_planner(
        llm=qwen_model,
        system_message=SystemMessage(content="""你是一个优秀的任务规划师。请按照以下要求工作：
1. 分析用户的任务需求，将其分解为具体的执行步骤
2. 每个步骤应该明确、可执行
3. 考虑步骤之间的依赖关系
4. 为每个步骤分配合适的工具
5. 确保最终能够完成用户的目标

可用工具:
- Search: 搜索最新信息
- Calculator: 数学计算
- Weather: 查询天气

请用中文进行规划。""")
    )

    # 创建执行器
    executor = load_agent_executor(
        llm=qwen_model,
        tools=tools,
        verbose=True,
        handle_parsing_errors=True
    )

    # 创建 PlanAndExecute 智能体
    agent = PlanAndExecute(
        planner=planner,
        executor=executor,
        verbose=True,
        handle_parsing_errors=True
    )

    return agent


def run_agent_example():
    """运行智能体示例"""

    print("初始化 Qwen-plus PlanAndExecute 智能体...")
    agent = create_qwen_plan_and_execute_agent()

    # 测试用例
    test_queries = [
        "请搜索人工智能的最新发展，然后计算未来5年市场规模的增长倍数，假设年增长率为15%",
        "查询北京的天气，然后计算如果温度下降5度后的温度是多少",
        "搜索机器学习的基本概念，然后解释其与人工智能的关系"
    ]

    for i, query in enumerate(test_queries, 1):
        print(f"\n{'=' * 50}")
        print(f"测试用例 {i}: {query}")
        print(f"{'=' * 50}")

        try:
            result = agent.invoke({"input": query})
            print(f"\n最终结果: {result}")
        except Exception as e:
            print(f"执行出错: {e}")


# 高级用法：自定义执行监控
class CustomPlanAndExecuteAgent:
    """自定义的 PlanAndExecute 智能体，增强监控功能"""

    def __init__(self):
        self.agent = create_qwen_plan_and_execute_agent()

    def run_with_monitoring(self, query):
        """带监控的执行"""
        print(f"开始处理查询: {query}")

        # 执行前的准备
        print("规划阶段开始...")

        # 执行查询
        result = self.agent.invoke({"input": query})

        print("任务执行完成!")
        return result


if __name__ == "__main__":
    # 运行示例
    run_agent_example()

    # 高级用法示例
    print("\n\n高级用法示例:")
    custom_agent = CustomPlanAndExecuteAgent()

    complex_query = """
    请执行以下复杂任务：
    1. 搜索2024年人工智能的主要趋势
    2. 计算如果这些趋势持续发展，到2030年AI行业的预计市场规模
    3. 分析这些趋势对就业市场的影响
    """

    result = custom_agent.run_with_monitoring(complex_query)
    print(f"复杂任务结果: {result}")